import socket
import struct

# QKD Packet Fields
HEADER_FORMAT = "!I I B B I I 32s"

class QKDAppPacket:
    def __init__(self, length, msg_id, encryption_algo, auth_algo, encryption_key_id, auth_key_id, auth_tag, payload):
        self.length = length
        self.msg_id = msg_id
        self.encryption_algo = encryption_algo
        self.auth_algo = auth_algo
        self.encryption_key_id = encryption_key_id
        self.auth_key_id = auth_key_id
        self.auth_tag = auth_tag
        self.payload = payload

    def pack(self):
        # Pack the header and payload into bytes
        header = struct.pack(HEADER_FORMAT, self.length, self.msg_id, self.encryption_algo, self.auth_algo,
                            self.encryption_key_id, self.auth_key_id, self.auth_tag)
        return header + self.payload

    @staticmethod
    def unpack(data):
        # Unpack the received data into the QKDAppPacket fields
        header_size = struct.calcsize(HEADER_FORMAT)
        header = data[:header_size]
        payload = data[header_size:]
        unpacked_data = struct.unpack(HEADER_FORMAT, header)
        return QKDAppPacket(*unpacked_data, payload)

class QKDAppTool:
    def __init__(self, bind_address=None):
        self.socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        if bind_address:
            self.socket.bind(bind_address)
            print(f"Server listening on {bind_address}")

    def send(self, address, packet):
        packed_data = packet.pack()
        self.socket.sendto(packed_data, address)
        print(f"Packet sent to {address}")

    def receive(self):
        data, addr = self.socket.recvfrom(4096)
        packet = QKDAppPacket.unpack(data)
        print(f"Packet received from {addr}:
 Length: {packet.length}
 Msg-ID: {packet.msg_id}
 Encryption Algorithm: {packet.encryption_algo}
 Authentication Algorithm: {packet.auth_algo}
 Encryption Key ID: {packet.encryption_key_id}
 Authentication Key ID: {packet.auth_key_id}
 Authentication Tag: {packet.auth_tag}
 Payload: {packet.payload.decode('utf-8')}")
        return packet, addr

    def close(self):
        self.socket.close()

if __name__ == "__main__":
    import threading
    import time

    def run_server():
        server = QKDAppTool(bind_address=('127.0.0.1', 12345))
        server.receive()
        server.close()

    def run_client():
        time.sleep(1)  # Ensure server is listening
        client = QKDAppTool()
        # Packet fields
        length = 100
        msg_id = 1
        encryption_algo = 1  # Example encryption algorithm ID
        auth_algo = 1  # Example authentication algorithm ID
        encryption_key_id = 101
        auth_key_id = 201
        auth_tag = b'auth_tag_example'.ljust(32, b'\0')  # Pad to 32 bytes
        payload = b'This is the protected payload'

        packet = QKDAppPacket(length, msg_id, encryption_algo, auth_algo, encryption_key_id, auth_key_id, auth_tag, payload)
        client.send(('127.0.0.1', 12345), packet)
        client.close()

    # Run server and client in separate threads for testing
    server_thread = threading.Thread(target=run_server)
    client_thread = threading.Thread(target=run_client)

    server_thread.start()
    client_thread.start()

    server_thread.join()
    client_thread.join()
